{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Advanced GRPO Fine-tuning for Mathematical Reasoning with Multi-Reward Training\n",
    "\n",
    "_Authored by: [Behrooz Azarkhalili](https://github.com/behroozazarkhalili)_\n",
    "\n",
    "This notebook demonstrates **advanced GRPO (Group Relative Policy Optimization)** for mathematical reasoning using a comprehensive multi-reward training system. We'll fine-tune a model on the GSM8K dataset with four specialized reward functions.\n",
    "\n",
    "**Key Features:**\n",
    "- **4 Reward Functions**: Format compliance, approximate matching, answer correctness, and number extraction\n",
    "- **Memory Efficient**: 4-bit quantization + LoRA for consumer GPUs\n",
    "- **Interactive Monitoring**: Real-time training metrics with trackio dashboard\n",
    "- **Structured Output**: Enforces step-by-step reasoning format\n",
    "\n",
    "The model learns to generate structured mathematical solutions with clear reasoning steps and accurate numerical answers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Installation and Setup\n",
    "\n",
    "Install the required packages for GRPO training with memory-efficient techniques."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install required packages for GRPO mathematical reasoning training\n",
    "!pip install transformers datasets trl bitsandbytes peft trackio"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GPU Environment Detection\n",
    "\n",
    "Verify GPU availability and display hardware specifications for optimal training configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "# Verify CUDA availability and display GPU specifications\n",
    "print(f\"CUDA available: {torch.cuda.is_available()}\")\n",
    "print(f\"Number of GPUs: {torch.cuda.device_count()}\")\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    # Display current GPU details for training optimization\n",
    "    print(f\"Current GPU: {torch.cuda.current_device()}\")\n",
    "    print(f\"GPU name: {torch.cuda.get_device_name()}\")\n",
    "    print(f\"GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")\n",
    "else:\n",
    "    # Provide guidance for enabling GPU in Colab\n",
    "    print(\"⚠️  No GPU available. This notebook requires a GPU for efficient training.\")\n",
    "    print(\"In Colab: Runtime → Change runtime type → Hardware accelerator → GPU\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Core Library Imports\n",
    "\n",
    "Import essential libraries for GRPO training, model configuration, and experiment tracking."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import trackio  # Experiment tracking dashboard\n",
    "import re       # Regex patterns for reward functions\n",
    "\n",
    "# GRPO training components\n",
    "from trl import GRPOConfig, GRPOTrainer\n",
    "\n",
    "# Model and tokenization\n",
    "from transformers import (\n",
    "    AutoModelForCausalLM,   # Causal language model loading\n",
    "    AutoTokenizer,          # Text tokenization\n",
    "    BitsAndBytesConfig,     # Quantization configuration\n",
    ")\n",
    "\n",
    "# Parameter-efficient fine-tuning\n",
    "from peft import LoraConfig, get_peft_model, TaskType\n",
    "\n",
    "# Dataset handling\n",
    "from datasets import load_dataset\n",
    "\n",
    "# Logging configuration\n",
    "import logging\n",
    "logging.basicConfig(level=logging.INFO)\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "# Suppress httpx request logs that appear during trackio usage\n",
    "logging.getLogger(\"httpx\").setLevel(logging.WARNING)\n",
    "logging.getLogger(\"gradio_client\").setLevel(logging.WARNING)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Selection and Configuration\n",
    "\n",
    "Choose a compact but capable model suitable for mathematical reasoning with memory constraints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Select model optimized for instruction-following and reasoning\n",
    "model_name = \"Qwen/Qwen2.5-3B-Instruct\"  # 3B parameter model balances capability and memory usage\n",
    "max_seq_length = 2048                     # Token limit for mathematical problems (reduce if OOM)\n",
    "\n",
    "print(f\"Loading model: {model_name}\")\n",
    "print(f\"Max sequence length: {max_seq_length}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure 4-bit quantization for ~75% memory reduction\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,                    # Enable 4-bit precision (vs 16-bit default)\n",
    "    bnb_4bit_quant_type=\"nf4\",           # NormalFloat4: optimal for neural network weights\n",
    "    bnb_4bit_compute_dtype=torch.float16, # Use FP16 for forward/backward passes\n",
    "    bnb_4bit_use_double_quant=True,      # Further quantize quantization constants\n",
    ")\n",
    "\n",
    "print(\"✅ 4-bit quantization configured\")\n",
    "print(\"   Memory reduction: ~75% vs FP16\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load model with quantization and automatic device mapping\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_name,\n",
    "    quantization_config=bnb_config,      # Apply 4-bit quantization\n",
    "    device_map=\"auto\",                   # Auto-distribute across available GPUs/CPU\n",
    "    trust_remote_code=True,              # Allow custom model code execution\n",
    "    torch_dtype=torch.float16,           # Use FP16 for non-quantized operations\n",
    ")\n",
    "\n",
    "# Load corresponding tokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "    model_name,\n",
    "    trust_remote_code=True               # Allow custom tokenizer code\n",
    ")\n",
    "\n",
    "# Ensure tokenizer has proper padding token for batch processing\n",
    "if tokenizer.pad_token is None:\n",
    "    tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "print(f\"✅ Model loaded successfully!\")\n",
    "print(f\"📊 Model parameters: ~{sum(p.numel() for p in model.parameters()) / 1e6:.1f}M\")\n",
    "print(f\"🧮 Quantized parameters: ~{sum(p.numel() for p in model.parameters() if hasattr(p, 'quant_type')) / 1e6:.1f}M\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LoRA Configuration\n",
    "\n",
    "Apply Low-Rank Adaptation to train only ~0.1% of parameters while maintaining performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure LoRA for mathematical reasoning adaptation\n",
    "lora_config = LoraConfig(\n",
    "    r=16,                              # Rank: adaptation capacity (16 good for reasoning tasks)\n",
    "    lora_alpha=32,                     # Scaling factor (typically 2x rank)\n",
    "    target_modules=[\"q_proj\", \"v_proj\"], # Focus on attention query/value for reasoning\n",
    "    lora_dropout=0.1,                  # Regularization to prevent overfitting\n",
    "    bias=\"none\",                       # Skip bias adaptation for simplicity\n",
    "    task_type=TaskType.CAUSAL_LM,      # Causal language modeling task\n",
    ")\n",
    "\n",
    "print(\"🔧 Applying LoRA adaptation to model...\")\n",
    "\n",
    "# Apply LoRA configuration to create trainable adapter\n",
    "model = get_peft_model(model, lora_config)\n",
    "\n",
    "# Display parameter efficiency\n",
    "print(\"📊 LoRA Training Parameters Summary:\")\n",
    "model.print_trainable_parameters()  # Shows trainable vs total parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GSM8K Dataset Setup\n",
    "\n",
    "Configure the GSM8K mathematical reasoning dataset with structured output format for step-by-step solutions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define structured output format for mathematical reasoning\n",
    "reasoning_start = \"<start_working_out>\"   # Begin reasoning section\n",
    "reasoning_end = \"<end_working_out>\"       # End reasoning section\n",
    "solution_start = \"<SOLUTION>\"            # Begin final answer\n",
    "solution_end = \"</SOLUTION>\"              # End final answer\n",
    "\n",
    "# System prompt that teaches the model our desired reasoning structure\n",
    "system_prompt = f\"\"\"You are a mathematical reasoning assistant.\n",
    "When given a math problem:\n",
    "1. Show your step-by-step work between {reasoning_start} and {reasoning_end}\n",
    "2. Provide your final numerical answer between {solution_start} and {solution_end}\n",
    "3. Be precise and show all calculation steps clearly.\"\"\"\n",
    "\n",
    "print(\"✅ Format tokens and system prompt defined\")\n",
    "print(f\"   Reasoning format: {reasoning_start} ... {reasoning_end}\")\n",
    "print(f\"   Solution format: {solution_start} ... {solution_end}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dataset processing utilities\n",
    "def extract_hash_answer(text):\n",
    "    \"\"\"Extract numerical answer from GSM8K format (#### marker)\"\"\"\n",
    "    if \"####\" not in text:\n",
    "        return None\n",
    "    # GSM8K uses format: \"Explanation... #### 42\"\n",
    "    return text.split(\"####\")[1].strip()\n",
    "\n",
    "def process_dataset_example(example):\n",
    "    \"\"\"Convert GSM8K example to conversation format for GRPO training\"\"\"\n",
    "    question = example[\"question\"]\n",
    "    answer = extract_hash_answer(example[\"answer\"])\n",
    "    \n",
    "    # Create conversation with system prompt for structured reasoning\n",
    "    prompt = [\n",
    "        {\"role\": \"system\", \"content\": system_prompt},\n",
    "        {\"role\": \"user\", \"content\": question},\n",
    "    ]\n",
    "    \n",
    "    return {\n",
    "        \"prompt\": prompt,           # Input conversation\n",
    "        \"answer\": answer,          # Ground truth for reward functions\n",
    "    }\n",
    "\n",
    "print(\"✅ Dataset processing functions defined\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load and preprocess GSM8K training dataset\n",
    "print(\"🔄 Loading GSM8K mathematical reasoning dataset...\")\n",
    "dataset = load_dataset(\"openai/gsm8k\", \"main\", split=\"train\")\n",
    "\n",
    "# Apply conversation formatting to all examples\n",
    "dataset = dataset.map(process_dataset_example)\n",
    "\n",
    "print(f\"✅ Dataset loaded and processed!\")\n",
    "print(f\"📊 Training examples: {len(dataset):,}\")\n",
    "print(f\"🎯 Sample question: {dataset[0]['prompt'][1]['content']}...\")\n",
    "print(f\"🎯 Sample answer: {dataset[0]['answer']}\")\n",
    "\n",
    "# Show structure of first example for verification\n",
    "print(f\"\\n📋 Example structure:\")\n",
    "print(f\"   Prompt: {len(dataset[0]['prompt'])} messages (system + user)\")\n",
    "print(f\"   Answer: {dataset[0]['answer']} (ground truth for rewards)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multi-Reward System Design\n",
    "\n",
    "Implement four complementary reward functions to evaluate different aspects of mathematical reasoning:\n",
    "1. **Exact Format Matching**: Perfect structure compliance  \n",
    "2. **Approximate Matching**: Partial credit for format elements\n",
    "3. **Answer Correctness**: Mathematical accuracy with graduated scoring\n",
    "4. **Number Extraction**: Ability to parse and output numerical results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compiled regex patterns for efficient reward computation\n",
    "match_format = re.compile(\n",
    "    rf\"^[\\s]{{0,}}\"                      # Optional whitespace at start\n",
    "    rf\"{reasoning_start}.+?{reasoning_end}.*?\"  # Reasoning section (non-greedy)\n",
    "    rf\"{solution_start}(.+?){solution_end}\"     # Solution section with capture group\n",
    "    rf\"[\\s]{{0,}}$\",                     # Optional whitespace at end\n",
    "    flags=re.MULTILINE | re.DOTALL       # Multi-line matching with . matching newlines\n",
    ")\n",
    "\n",
    "match_numbers = re.compile(\n",
    "    rf\"{solution_start}.*?([\\d\\.]{{1,}})\", # Extract numbers from solution section\n",
    "    flags=re.MULTILINE | re.DOTALL        # Flexible pattern matching\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reward Function 1: Exact Format Compliance\n",
    "def match_format_exactly(completions, **kwargs):\n",
    "    \"\"\"\n",
    "    High reward (3.0) for perfect format adherence\n",
    "    Ensures model learns the complete structured output pattern\n",
    "    \"\"\"\n",
    "    scores = []\n",
    "    for completion in completions:\n",
    "        response = completion[0][\"content\"]\n",
    "        # Check if response matches complete format pattern\n",
    "        score = 3.0 if match_format.search(response) is not None else 0.0\n",
    "        scores.append(score)\n",
    "    return scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reward Function 2: Partial Format Credit\n",
    "def match_format_approximately(completions, **kwargs):\n",
    "    \"\"\"\n",
    "    Graduated scoring for format elements\n",
    "    Encourages learning individual components even if not perfect\n",
    "    \"\"\"\n",
    "    scores = []\n",
    "    for completion in completions:\n",
    "        response = completion[0][\"content\"]\n",
    "        score = 0\n",
    "        \n",
    "        # Award +0.5 for correct token count, -0.5 for wrong count\n",
    "        score += 0.5 if response.count(reasoning_start) == 1 else -0.5\n",
    "        score += 0.5 if response.count(reasoning_end) == 1 else -0.5\n",
    "        score += 0.5 if response.count(solution_start) == 1 else -0.5\n",
    "        score += 0.5 if response.count(solution_end) == 1 else -0.5\n",
    "        \n",
    "        scores.append(score)\n",
    "    return scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reward Function 3: Mathematical Accuracy\n",
    "def check_answer_correctness(prompts, completions, answer, **kwargs):\n",
    "    \"\"\"\n",
    "    Graduated scoring for mathematical accuracy:\n",
    "    - 3.0: Exact match\n",
    "    - 1.5: Within 10% (close answer)\n",
    "    - 0.5: Within 20% (reasonable attempt)\n",
    "    - -0.5: Wrong answer (penalty for incorrect math)\n",
    "    \"\"\"\n",
    "    responses = [completion[0][\"content\"] for completion in completions]\n",
    "    \n",
    "    # Extract answers using format pattern\n",
    "    extracted_responses = [\n",
    "        guess.group(1) if (guess := match_format.search(r)) is not None else None\n",
    "        for r in responses\n",
    "    ]\n",
    "    \n",
    "    scores = []\n",
    "    for guess, true_answer in zip(extracted_responses, answer):\n",
    "        if guess is None:  # No extractable answer\n",
    "            scores.append(0)\n",
    "            continue\n",
    "            \n",
    "        # Exact string match gets full points\n",
    "        if guess.strip() == true_answer.strip():\n",
    "            scores.append(3.0)\n",
    "        else:\n",
    "            # Try numerical comparison for partial credit\n",
    "            try:\n",
    "                ratio = float(guess) / float(true_answer)\n",
    "                if 0.9 <= ratio <= 1.1:      # Within 10%\n",
    "                    scores.append(1.5)\n",
    "                elif 0.8 <= ratio <= 1.2:    # Within 20%\n",
    "                    scores.append(0.5)\n",
    "                else:                         # Wrong answer\n",
    "                    scores.append(-0.5)\n",
    "            except (ValueError, ZeroDivisionError):\n",
    "                scores.append(-0.5)           # Invalid numerical format\n",
    "    \n",
    "    return scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reward Function 4: Number Extraction Ability  \n",
    "def check_numbers_extraction(prompts, completions, answer, **kwargs):\n",
    "    \"\"\"\n",
    "    Tests the model's ability to extract numerical values from solution sections\n",
    "    Complementary to exact format matching - focuses on parsing capability\n",
    "    \"\"\"\n",
    "    responses = [completion[0][\"content\"] for completion in completions]\n",
    "    \n",
    "    # Extract numbers from solution sections using number pattern\n",
    "    extracted_responses = [\n",
    "        guess.group(1) if (guess := match_numbers.search(r)) is not None else None\n",
    "        for r in responses\n",
    "    ]\n",
    "    \n",
    "    scores = []\n",
    "    for guess, true_answer in zip(extracted_responses, answer):\n",
    "        if guess is None:  # No extractable number\n",
    "            scores.append(0)\n",
    "            continue\n",
    "            \n",
    "        try:\n",
    "            # Simple numerical equality check\n",
    "            true_val = float(true_answer.strip())\n",
    "            guess_val = float(guess.strip())\n",
    "            # Binary scoring: correct (1.5) or incorrect (0)\n",
    "            scores.append(1.5 if guess_val == true_val else 0.0)\n",
    "        except (ValueError, TypeError):\n",
    "            scores.append(0)  # Invalid number format\n",
    "    \n",
    "    return scores"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GRPO Training Setup\n",
    "\n",
    "Configure training parameters optimized for mathematical reasoning with memory constraints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure GRPO training parameters for mathematical reasoning\n",
    "training_args = GRPOConfig(\n",
    "    # Learning parameters optimized for reasoning tasks\n",
    "    learning_rate=5e-6,              # Conservative LR to prevent destabilizing reasoning\n",
    "    \n",
    "    # Memory-efficient batch configuration\n",
    "    per_device_train_batch_size=2,   # Small batch for GPU memory constraints\n",
    "    gradient_accumulation_steps=8,   # Effective batch size = 2 * 8 = 16\n",
    "    \n",
    "    # Sequence length limits for mathematical problems\n",
    "    max_prompt_length=1024,          # Sufficient for complex word problems\n",
    "    max_completion_length=1024,      # Room for detailed step-by-step reasoning\n",
    "    \n",
    "    # Training duration and monitoring\n",
    "    max_steps=10,                    # Short demo run (increase to 500+ for production)\n",
    "    logging_steps=1,                 # Log metrics every step for close monitoring\n",
    "    \n",
    "    # Stability and output configuration\n",
    "    output_dir=\"./trl_grpo_outputs\",\n",
    "    max_grad_norm=0.1,               # Aggressive gradient clipping for stable training\n",
    "    report_to=\"trackio\",                # use trackio for experiment tracking (instead of wandb/tensorboard)\n",
    ")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create unique run name with timestamp to ensure fresh tracking\n",
    "import datetime\n",
    "timestamp = datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
    "run_name = f\"qwen2.5-3b-gsm8k-grpo-{timestamp}\"\n",
    "\n",
    "# Initialize trackio experiment tracking with unique run name\n",
    "trackio.init(\n",
    "    project=\"GRPO-Mathematical-Reasoning\",  # Project name for organization\n",
    "    name=run_name,                         # Unique run identifier with timestamp\n",
    "    config={\n",
    "        # Model and dataset configuration\n",
    "        \"model_name\": \"Qwen/Qwen2.5-3B-Instruct\",\n",
    "        \"dataset\": \"GSM8K\", \n",
    "        \"technique\": \"GRPO + LoRA + 4-bit\",\n",
    "        \n",
    "        # Training hyperparameters\n",
    "        \"learning_rate\": training_args.learning_rate,\n",
    "        \"batch_size\": training_args.per_device_train_batch_size,\n",
    "        \"gradient_accumulation_steps\": training_args.gradient_accumulation_steps,\n",
    "        \"effective_batch_size\": training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps,\n",
    "        \"max_steps\": training_args.max_steps,\n",
    "        \n",
    "        # LoRA configuration\n",
    "        \"lora_r\": 16,\n",
    "        \"lora_alpha\": 32,\n",
    "        \n",
    "        # GRPO-specific settings\n",
    "        \"num_generations\": training_args.num_generations,  # Default: 8 generations per step\n",
    "        \"max_prompt_length\": training_args.max_prompt_length,\n",
    "        \"max_completion_length\": training_args.max_completion_length,\n",
    "        \n",
    "        # Reward system\n",
    "        \"num_reward_functions\": 4,\n",
    "    }\n",
    ")\n",
    "\n",
    "print(\"🎯 GRPO Configuration Summary:\")\n",
    "print(f\"   Learning rate: {training_args.learning_rate}\")\n",
    "print(f\"   Effective batch size: {training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps}\")\n",
    "print(f\"   Training steps: {training_args.max_steps}\")\n",
    "print(f\"   Generations per step: {training_args.num_generations}\")\n",
    "print(f\"✅ Trackio experiment tracking initialized\")\n",
    "print(f\"📊 Run name: {run_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Trainer Initialization with Trackio Integration\n",
    "\n",
    "Set up the GRPO trainer with our multi-reward system and experiment tracking."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize GRPO trainer with multi-reward system\n",
    "# trackio_callback = TrackioCallback()  # Create trackio logging callback\n",
    "\n",
    "trainer = GRPOTrainer(\n",
    "    model=model,                      # LoRA-adapted quantized model\n",
    "    reward_funcs=[                    # Four complementary reward functions\n",
    "        match_format_exactly,         # Perfect structure compliance\n",
    "        match_format_approximately,   # Partial format credit\n",
    "        check_answer_correctness,     # Mathematical accuracy\n",
    "        check_numbers_extraction,     # Number parsing ability\n",
    "    ],\n",
    "    args=training_args,               # Training configuration\n",
    "    train_dataset=dataset,            # Processed GSM8K dataset\n",
    ")\n",
    "\n",
    "print(\"✅ GRPO Trainer initialized successfully!\")\n",
    "print(f\"📊 Training dataset: {len(dataset):,} examples\")\n",
    "print(f\"🎯 Reward functions: {len(trainer.reward_funcs)} active\")\n",
    "print(f\"📈 Trackio integration: Enabled\")\n",
    "print(f\"🔄 Ready for training with {training_args.num_generations} generations per step\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Begin GRPO Training\n",
    "\n",
    "Start the training process with real-time reward monitoring. Watch for gradual improvement in both format compliance and mathematical accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Execute GRPO training with multi-reward optimization\n",
    "print(\"🚀 Starting GRPO training...\")\n",
    "print(\"📊 Monitor metrics: reward scores, KL divergence, policy gradients\")\n",
    "print(\"🔍 Trackio will log: losses, rewards, learning rate, gradients\")\n",
    "\n",
    "# Run the training process\n",
    "trainer.train()\n",
    "\n",
    "# Complete the trackio experiment\n",
    "trackio.finish()\n",
    "\n",
    "print(\"✅ Training completed successfully!\")\n",
    "print(f\"💾 Model saved to: {training_args.output_dir}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Experiment Dashboard\n",
    "\n",
    "Launch the interactive trackio dashboard to analyze training progress, reward evolution, and model performance metrics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Launch interactive trackio dashboard for experiment analysis\n",
    "# View training curves, reward progression, loss evolution, and hyperparameter effects\n",
    "trackio.show(project=\"GRPO-Mathematical-Reasoning\")\n",
    "\n",
    "\n",
    "# Alternative: Launch from command line with: trackio show --project \"GRPO-Mathematical-Reasoning\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Evaluation and Testing\n",
    "\n",
    "Test the trained model's mathematical reasoning capability with structured output validation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define model testing function with optimized generation parameters\n",
    "def test_model(question, max_length=512):\n",
    "    \"\"\"\n",
    "    Test the trained model on mathematical questions\n",
    "    \n",
    "    Args:\n",
    "        question (str): Mathematical problem to solve\n",
    "        max_length (int): Maximum tokens to generate\n",
    "        \n",
    "    Returns:\n",
    "        str: Model's structured response with reasoning and solution\n",
    "    \"\"\"\n",
    "    # Format input using conversation template\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": system_prompt},\n",
    "        {\"role\": \"user\", \"content\": question},\n",
    "    ]\n",
    "    \n",
    "    # Apply chat template and tokenize\n",
    "    text = tokenizer.apply_chat_template(\n",
    "        messages,\n",
    "        add_generation_prompt=True,      # Add assistant prompt\n",
    "        tokenize=False,                  # Return string, not tokens\n",
    "    )\n",
    "    \n",
    "    # Tokenize and move to appropriate device\n",
    "    inputs = tokenizer(text, return_tensors=\"pt\").to(model.device)\n",
    "    \n",
    "    print(f\"🤔 Processing: {question}\")\n",
    "    \n",
    "    # Generate response with reasoning-optimized parameters\n",
    "    with torch.no_grad():\n",
    "        outputs = model.generate(\n",
    "            **inputs,\n",
    "            max_new_tokens=max_length,\n",
    "            temperature=0.7,                    # Balance creativity and consistency\n",
    "            do_sample=True,                     # Enable sampling for varied reasoning paths\n",
    "            top_p=0.9,                         # Nucleus sampling for quality\n",
    "            pad_token_id=tokenizer.eos_token_id,\n",
    "            repetition_penalty=1.1,            # Reduce repetitive reasoning steps\n",
    "            length_penalty=1.0,                # Neutral preference for response length\n",
    "            early_stopping=True,               # Stop at natural completion\n",
    "        )\n",
    "    \n",
    "    # Decode and extract only the generated portion\n",
    "    response = tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
    "    generated_text = response[len(text):].strip()\n",
    "    \n",
    "    return generated_text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test model on GSM8K problem\n",
    "gsm8k_question = \"Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?\"\n",
    "expected_answer = \"72\"\n",
    "\n",
    "# Generate response\n",
    "gsm8k_response = test_model(gsm8k_question, max_length=768)\n",
    "\n",
    "print(f\"Question: {gsm8k_question}\")\n",
    "print(f\"Model Response:\\n{gsm8k_response}\")\n",
    "\n",
    "# Validate format compliance\n",
    "has_reasoning = reasoning_start in gsm8k_response and reasoning_end in gsm8k_response\n",
    "has_solution = solution_start in gsm8k_response and solution_end in gsm8k_response\n",
    "\n",
    "print(f\"\\nFormat Check:\")\n",
    "print(f\"Reasoning section: {has_reasoning}\")\n",
    "print(f\"Solution section: {has_solution}\")\n",
    "\n",
    "# Check answer accuracy if solution section exists\n",
    "if has_solution:\n",
    "    try:\n",
    "        solution_text = gsm8k_response.split(solution_start)[1].split(solution_end)[0].strip()\n",
    "        extracted_number = ''.join(filter(str.isdigit, solution_text))\n",
    "        expected_number = ''.join(filter(str.isdigit, expected_answer))\n",
    "        is_correct = extracted_number == expected_number\n",
    "        \n",
    "        print(f\"Extracted: {solution_text}\")\n",
    "        print(f\"Expected: {expected_answer}\")\n",
    "        print(f\"Correct: {is_correct}\")\n",
    "    except:\n",
    "        print(\"Could not extract solution\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Clean Up Resources\n",
    "\n",
    "Free GPU memory and clear cached tensors for optimal resource management."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "def remove_trackio_project(project_name):\n",
    "    \"\"\"Remove a trackio project by deleting its database file\"\"\"\n",
    "    cache_dir = Path.home() / \".cache\" / \"huggingface\" / \"trackio\"\n",
    "    db_file = cache_dir / f\"{project_name}.db\"\n",
    "    \n",
    "    if db_file.exists():\n",
    "        db_file.unlink()\n",
    "        print(f\"Removed trackio project: {project_name}\")\n",
    "    else:\n",
    "        print(f\"Project not found: {project_name}\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clean up trackio experiment database to free storage space\n",
    "# WARNING: This permanently deletes all experiment logs and metrics\n",
    "remove_trackio_project(\"GRPO-Mathematical-Reasoning\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Free GPU memory and clear Python garbage collection\n",
    "import gc\n",
    "\n",
    "torch.cuda.empty_cache()  # Clear PyTorch CUDA memory cache\n",
    "gc.collect()              # Run Python garbage collector\n",
    "\n",
    "print(\"✅ GPU memory cache cleared\")\n",
    "print(\"✅ Python garbage collection completed\")\n",
    "print(\"🧹 Resources freed for other processes\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "### Papers and Research\n",
    "- **GRPO Algorithm**: [Group Relative Policy Optimization](https://arxiv.org/abs/2402.03300) - The original GRPO paper introducing group-based relative policy optimization\n",
    "- **GSM8K Dataset**: [Training Verifiers to Solve Math Word Problems](https://arxiv.org/abs/2110.14168) - Cobbe et al., OpenAI\n",
    "- **LoRA**: [Low-Rank Adaptation of Large Language Models](https://arxiv.org/abs/2106.09685) - Hu et al., Microsoft\n",
    "- **QLoRA**: [Efficient Finetuning of Quantized LLMs](https://arxiv.org/abs/2305.14314) - Dettmers et al., 4-bit quantization for efficient training\n",
    "\n",
    "### Libraries and Frameworks\n",
    "- **TRL (Transformers Reinforcement Learning)**: [HuggingFace TRL](https://github.com/huggingface/trl) - Official library for RLHF and advanced training techniques\n",
    "- **Transformers**: [HuggingFace Transformers](https://github.com/huggingface/transformers) - State-of-the-art NLP library\n",
    "- **PEFT**: [Parameter-Efficient Fine-Tuning](https://github.com/huggingface/peft) - Efficient adaptation methods\n",
    "- **BitsAndBytes**: [8-bit & 4-bit Quantization](https://github.com/TimDettmers/bitsandbytes) - Memory-efficient training\n",
    "\n",
    "### Models Used\n",
    "- **Qwen2.5-3B-Instruct**: [Qwen Model Series](https://github.com/QwenLM/Qwen2.5) - Alibaba's instruction-tuned language model\n",
    "- **Alternative Models**: Gemma-2B, DialoGPT, GPT-2 (configurable in the notebook)\n",
    "\n",
    "### Datasets\n",
    "- **GSM8K**: [OpenAI GSM8K](https://huggingface.co/datasets/openai/gsm8k) - Grade School Math 8K problems dataset\n",
    "- **Format**: Mathematical word problems requiring multi-step reasoning and numerical answers\n",
    "\n",
    "### Key Concepts\n",
    "- **Reinforcement Learning from Human Feedback (RLHF)**: Training language models using reward signals\n",
    "- **Group Relative Policy Optimization**: Advanced RL technique comparing responses in groups rather than absolute scoring\n",
    "- **Structured Generation**: Teaching models to follow specific output formats with reasoning sections\n",
    "- **Multi-Reward Training**: Using multiple reward functions for comprehensive evaluation"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "behrooz",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
